import os, argparse
from sys import path
import cv2
import numpy as np
import multiprocessing as mp
import math
import atexit
import time
from attack import WSA
from utils.f_measure import f_measure
from utils.weighted_l2_loss import weighted_l2_loss
from utils.dataloader import data_filter, dataloader
from utils.modelloader import load_model
from utils.path_check import path_check

os.environ['GLOG_minloglevel'] = '3'

parser = argparse.ArgumentParser(description='WSA Attack')
parser.add_argument('--sourcedir', default='clean imgs dir')
parser.add_argument('--gtdir', default='gt dir')
parser.add_argument('--targetdir', default='adv samples dir for saving')
parser.add_argument('--skeletondir', default='skeleton dir for saving')
parser.add_argument('--progressdir', default='attack progress dir for saving')
parser.add_argument('--eps', type=int, default=16,  help ='pertrbation budget')
parser.add_argument('--iters', type=int, default=100, help='Number of Iterations')
parser.add_argument('--dataset', default='sklarge', help='dataset name')
parser.add_argument('--modelname', default='deepflux', help='model name')
parser.add_argument('--gpu', type=int, default=0, help='cuda index')
parser.add_argument('--attacker', default='wsa', help='attack method')
parser.add_argument('--num', type=int, default=1, help='number of process')
args = parser.parse_args()

path_check(args.targetdir)
path_check(args.skeletondir)
path_check(args.progressdir)

def main(datalist, cuda=0, start=0):
    net = load_model(args.modelname, args.dataset, cuda)

    for index in range(len(datalist)):
        image, gt, name = datalist[index]

        target_img_path = os.path.join(args.targetdir, name + '.png')
        target_skl_path = os.path.join(args.skeletondir, name + '.png')
        progress_path = os.path.join(args.progressdir, name + '.png')

        if args.attacker == 'origin':
            # skeleton test
            skeleton = net(image)
            # score = f_measure(skeleton, gt)
            score = .0

            skl_to_save = (255 * skeleton).round().astype(np.uint8)
            cv2.imwrite(target_skl_path, skl_to_save)

        elif args.attacker == 'wsa' or args.attacker == 'frequency':
            # WSA
            adv, skeleton, score, params = WSA(image, net, weighted_l2_loss, eps=args.eps, gt=(gt if args.attacker == 'frequency' else None), show=False, save_dir=progress_path)
            
            adv_to_save = (255 * adv).round().astype(np.uint8)
            skl_to_save = (255 * skeleton).round().astype(np.uint8)
            cv2.imwrite(target_img_path, adv_to_save)
            cv2.imwrite(target_skl_path, skl_to_save)
        
        else:
            exit()

        print("no." + str(start + index + 1) + " " + str(round(100*(index+1)/len(datalist), 1)) + "% measure: " + str(round(score, 4)))
        if not args.attacker == 'origin':
            print('adv: ' + target_img_path)
        print('skl: ' + target_skl_path)
        # print('prg: ' + progress_path)
    print('done')


pools = []
def multi_attack(num, gpu_pools=[0, 1, 2]):
    datalist = dataloader(args.sourcedir, args.gtdir, args.dataset)
    datalist = data_filter(datalist, args.skeletondir)

    if num == 1:
        main(datalist, args.gpu, 0)
        return

    sub_part_num = int(math.ceil(len(datalist) / num)) + 1
    
    for i in range(num):
        cuda = i % len(gpu_pools)
        start = i * sub_part_num
        stop = min(len(datalist), start + sub_part_num)

        p = mp.Process(target=main, args=(datalist[start: stop], cuda, start, ))
        pools.append(p)

    for p in pools:
        p.start()

    while True:
        time.sleep(3)
        flag = True
        for p in pools:
            if p.is_alive():
                flag = False
        if flag:
            break


@atexit.register
def clean():
    for p in pools:
        if p.is_alive():
            print('kill ' + str(p.pid))
            p.terminate()





if __name__ == '__main__':
    if args.modelname == 'deepflux':
        num = 12
    elif args.modelname == 'fsds':
        num = 12
    elif args.modelname == 'hifi':
        num = 3
    else:
        num = 1
    multi_attack(num)